#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>

#include "config.h"
#include "libxsp.h"

#include "psd_conn.h"
#include "psd_modules.h"
#include "psd_measurement.h"
#include "psd_conn_tcp.h"
#include "psd_conn_file.h"
#include "psd_config.h"
#include "psd_logger.h"
#include "psd_tpool.h"

#include "compat.h"

void *psd_lsrc_accept_handler(void *arg);
void *psd_lsrc_handle_client(void *arg);
int psd_lsrc_init();
void psd_lsrc_read_config();

static int *sockets;
static int num_sockets;

static psdModule psd_lsrc_module = {
	.desc = "XSP Depot Source Module",
	.dependencies = "tcp",
	.init = psd_lsrc_init
};

struct psd_lsrc_config_t {
	int port;
	int send_bufsize;
	int send_timeout;
	char *filename;
	int output_amt;
};

static struct psd_lsrc_config_t psdLsrcConfig = {
	.port =         5002,
	.send_bufsize = 0,
	.send_timeout = 0,
	.filename  = "/dev/zero",
	.output_amt   = 8388608
};

void psd_lsrc_read_config() {
	int val;
	char *str_val;

	if (psd_depot_settings_get_int("lsrc", "output_amt", &val) == 0) {
		psdLsrcConfig.output_amt = val;
	}

	if (psd_depot_settings_get("lsrc", "filename", &str_val) == 0) {
		psdLsrcConfig.filename = str_val;
	}

	if (psd_depot_settings_get_int("lsrc", "port", &val) == 0) {
		psdLsrcConfig.port = val;
	}

	if (psd_depot_settings_get_int("lsrc", "send_bufsize", &val) == 0) {
		psdLsrcConfig.send_bufsize = val;
	}

	if (psd_depot_settings_get_int("lsrc", "send_timeout", &val) == 0) {
		psdLsrcConfig.send_timeout = val;
	}

}

psdModule *module_info() {
	return &psd_lsrc_module;
}

int psd_lsrc_init() {
	struct addrinfo *ai;
	int i;

	psd_lsrc_read_config();

	sockets = listen_port(IPPROTO_TCP, PF_UNSPEC, psdLsrcConfig.port, &num_sockets, &ai);
	if (!sockets) {
		psd_err(0, "couldn't listen on port: %d", psdLsrcConfig.port);
		goto error_exit;
	}

	freeaddrinfo(ai);

	for(i = 0; i < num_sockets; i++) {

		if (psdLsrcConfig.send_bufsize > 0) {
			setsockopt(sockets[i], SOL_SOCKET, SO_SNDBUF,
					&(psdLsrcConfig.send_bufsize),
					sizeof(int));
		}

	}

	if (psd_tpool_exec(psd_lsrc_accept_handler, NULL) != 0) {
		psd_err(0, "couldn't execute socket handler thread");
		goto error_exit2;
	}

	return 0;

error_exit2:
	for(i = 0; i < num_sockets; i++)
		close(sockets[i]);
error_exit:
	return -1;
}

void *psd_lsrc_accept_handler(void *arg) {
	SOCKLEN_T peer_sa_len;
	fd_set socket_set;
	int n, i, high;
	struct sockaddr_storage sa;
	int cfd;

	while(1) {

		FD_ZERO(&socket_set);

		high = 0;

		for(i = 0; i < num_sockets; i++) {
			FD_SET(sockets[i], &socket_set);
			if (sockets[i] > high)
				high = sockets[i];
		}

		n = select(high + 1, &socket_set, NULL, NULL, NULL);

		for(i = 0; i < num_sockets; i++) {
			if (FD_ISSET(sockets[i], &socket_set)) {
				psdConn *conn;

				peer_sa_len = sizeof(struct sockaddr_storage);

				cfd = accept(sockets[i], (struct sockaddr *) &sa, &peer_sa_len);
				if (cfd < 0) {
					perror("main_loop: accept failed");
					continue;
				}

				conn = psd_conn_tcp_alloc(cfd, 0);
				if (!conn) {
					close(cfd);
					continue;
				}

				psd_tpool_exec(psd_lsrc_handle_client, conn);
			}
		}
	}

	return NULL;
}

void* psd_lsrc_handle_client(void *arg) {
	int file_fd;
	uint64_t amt = 0;
	uint64_t amt_remaining;
	char *buf;
	int bufsize;
	psdConn *conn = arg;
	psdConn *file_conn;

	psd_info(1, "received connection from %s", conn->description);

	if (psdLsrcConfig.send_bufsize == 0) {
		bufsize = 8*1024*1024;
	} else {
		bufsize = psdLsrcConfig.send_bufsize;
	}

	psd_info(5, "buffer size set to %d", bufsize);

	buf = malloc(sizeof(char) * bufsize);
	if (!buf) {
		psd_err(1, "failed to allocate bufferspace for %s", conn->description);
		goto error_exit;
	}

	file_fd = open(psdLsrcConfig.filename, O_RDONLY, 0);
	if (file_fd < 0) {
		psd_err(1, "failed to open %s", psdLsrcConfig.filename);
		goto error_exit2;
	}

	file_conn = psd_alloc_conn_file(file_fd);
	if (!file_conn) {
		psd_err(1, "failed to allocate file connection for %s", conn->description);
		goto error_exit3;
	}

	if (psdLsrcConfig.output_amt == 0) {
		amt_remaining = (uint64_t) -1; // "INFINITE"
	} else {
		amt_remaining = psdLsrcConfig.output_amt;
	}

	psd_info(5, "buffer size set to %d", bufsize);

	psd_measure_conn(conn);

	while(amt_remaining > 0) {
		int amt_to_send;
		int n;
		int amt_wrote;

		n = psd_conn_read(file_conn, buf, bufsize, 0);
		if (n <= 0) {
			goto out;
		}

		psd_info(10, "read in %d bytes", n);

		amt_to_send = n;
		if (amt_remaining < amt_to_send) {
			amt_to_send = amt_remaining;
			psd_info(10, "only sending %d of the %d bytes read", amt_to_send, n);
		}

		amt_remaining -= amt_to_send;

		amt_wrote = 0;
		while(amt_to_send > 0) {
			n = psd_conn_write(conn, buf + amt_wrote, amt_to_send, 0);
			if (n <= 0) {
				goto out;
			}

			amt_to_send -= n;
			amt_wrote += n;
			amt += n;
			psd_info(10, "wrote %d bytes to %s", n, conn->description);
		}

	}
out:
	psd_info(1, "Connection to %s closed. Wrote %llu bytes.",
			conn->description, amt);

	free(buf);
	psd_conn_shutdown(file_conn, PSD_SEND_SIDE | PSD_RECV_SIDE);
	psd_conn_free(file_conn);
	psd_stop_measure_conn(conn);
	psd_conn_shutdown(conn, PSD_SEND_SIDE | PSD_RECV_SIDE);
	psd_conn_free(conn);

	return NULL;

error_exit3:
	close(file_fd);
error_exit2:
	free(buf);
error_exit:
	psd_conn_shutdown(conn, PSD_SEND_SIDE | PSD_RECV_SIDE);
	psd_conn_free(conn);
	return NULL;
}
